# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/8 下午2:46
# @Author  : shutao
# @FileName: QTM_circuit_synthesis.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import numpy as np
import paddle
import os

from paddle_quantum.circuit import UAnsatz


class Circuit(UAnsatz):

    def gate_fidelity(self, target):
        f = paddle.matmul(self.U, target, transpose_y=True)
        trace = paddle.trace(f)
        return paddle.abs(trace) / 2**self.n

    @property
    def complexity(self):
        c = 0
        for g, _, _ in self._UAnsatz__history:
            if g == 'u':
                c += 1
            elif g == 'CNOT':
                c += 8
        return c

    def output(self, path):
        history = []
        for idx, (g, w, p) in enumerate(self._UAnsatz__history):
            g = self._UAnsatz__history[idx][g]
            w = self._UAnsatz__history[idx][w]
            p = self._UAnsatz__history[idx][p]
            if g == 'u' or g == 'ry':
                theta = self._UAnsatz__param[p[0]].numpy()[0]
                history.append('R %d %f\n' % (w[0], np.mod(theta, 2 * np.pi)))
            elif g == 'CNOT':
                history.append('C %d %d\n' % tuple(w))
            else:
                raise OSError('未知的门`%s`' % g)
        with open(path, 'w') as f:
            for i in history:
                f.write(i)


class Net(paddle.nn.Layer):
    def __init__(self, n, p, circuit, dtype="float64"):
        super(Net, self).__init__()
        self.n = n
        target = 1 / 2 ** 0.5 * np.array([[1, -1], [1, 1]])  # 目标酉矩阵
        self.target_tensor = paddle.to_tensor(target)
        self.theta = self.create_parameter(shape=[p],
                                           default_initializer=paddle.nn.initializer.Uniform(low=-np.pi/2, high=np.pi/2),
                                           dtype=dtype, is_bias=False)
        self.circuit = circuit

    def forward(self):
        cir = self.circuit(self.theta)
        loss = 1 - cir.gate_fidelity(self.target_tensor)
        return loss, cir


class Optimizer:
    def __init__(self, net, learning_rate):
        scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=learning_rate, T_max=100, eta_min=1e-3,
                                                             last_epoch=99, verbose=False)
        opt = paddle.optimizer.Adam(learning_rate=scheduler, parameters=net.parameters())
        self.net = net
        self.scheduler = scheduler
        self.opt = opt

    def train(self, max_iters=500, log_iter=50, threshold=1e-5, display=True):
        last = 1
        for itr in range(1, max_iters + 1):
            loss, cir = self.net()
            loss.backward()
            self.opt.minimize(loss)
            self.opt.clear_grad()
            self.scheduler.step()
            if itr % log_iter == 0:
                l = loss.numpy()[0]
                if display:
                    print("iter:", "%3d" % itr, "  loss:", "%.8f" % (last-l),
                          "  theta:", "%.6f" % self.net.theta.numpy()[0])
                if abs(l - last) < threshold:
                    break
                last = l
        return last


def run_main():
    n = 1
    LR = 0.1
    ITR = 1000
    # outer scope
    def circuit(theta):
        cir = Circuit(1)
        cir.ry(theta[0], 0)
        return cir

    net = Net(n, 1, circuit)  # 定义网络
    opt = Optimizer(net, LR)  # 定义自动优化器
    opt.train(ITR, threshold=1e-15)  # 优化参数

    theta_opt = net.theta.numpy()[0]  # 优化好的参数
    print("优化后的参数 theta: ", "%.4f" % theta_opt)

    cir = circuit(net.theta)  # 将优化好的参数加载到量子电路中
    cir.output('Question_1_Answer.txt')  # 输出结果到文件


if __name__ == '__main__':
    run_main()
